#!/usr/bin/python
# -*-coding:utf-8-*-
import os
import pandas as pd
import numpy as np

## TODO - for test
test_data_dir = './zbc_gplearn_factor_mining/raw_X_data'

test_data = pd.read_hdf(os.path.join(test_data_dir, 'X_user_behavior_data_v1.h5'))


df = test_data.copy()
random_idx = np.arange(0, df.shape[0])
np.random.shuffle(random_idx)
df = df.iloc[random_idx, :]
Xname = 'reg_get_five_year_del_user_optional_num'

'''横截面操作'''
def cs_rank_func(df, Xname):
    Xnewname = 'cs_rank_%s' % (Xname)

    try:
        df = df.copy()

        Xdata = df.groupby('date',
                           as_index=False,
                           sort=True,
                           group_keys=False)[[Xname]].rank(na_option='keep', ascending=True, pct=True)

        df[Xnewname] = Xdata[Xname]
        flag = True
    except:
        print(Xname, 'do cs_rank_func exception!')
        flag = False

    return df, flag

def cs_zscroe_func(df, Xname):
    Xnewname = 'cs_zscore_%s' % (Xname)

    def zscore(data, Xname):
        data = data[[Xname]].copy()

        data[Xname] = (data[Xname] - data[Xname].mean(skipna=True)) / data[Xname].std(skipna=True)

        return data

    try:
        df = df.copy()

        Xdata = df.groupby('date',
                           as_index=False,
                           sort=True,
                           group_keys=False).apply(zscore, Xname)

        df[Xnewname] = Xdata[Xname]
        flag = True
    except:
        print(Xname, 'do cs_rank_func exception!')
        flag = False

    return df, flag

def cs_demean_func(df, Xname):
    Xnewname = 'cs_demean_%s' % (Xname)

    def demean_opt(data, Xname):
        data = data[[Xname]].copy()

        data[Xname] = data[Xname] - data[Xname].mean(skipna=True)

        return data

    try:
        df = df.copy()

        Xdata = df.groupby('date',
                           as_index=False,
                           sort=True,
                           group_keys=False).apply(demean_opt, Xname)

        df[Xnewname] = Xdata[Xname]
        flag = True
    except:
        print(Xname, 'do cs_demean_func exception!')
        flag = False

    return df, flag


'''时间序列'''
def ts_argmax_func(df, Xname):
    Xnewname = 'ts_argmax_%s' % (Xname)

    def demean_opt(data, Xname):
        data = data[[Xname]].copy()

        data[Xname] = data[Xname] - data[Xname].mean(skipna=True)

        return data

    try:
        df = df.copy()

        Xdata = df.groupby('date',
                           as_index=False,
                           sort=True,
                           group_keys=False).apply(demean_opt, Xname)

        df[Xnewname] = Xdata[Xname]
        flag = True
    except:
        print(Xname, 'do cs_demean_func exception!')
        flag = False

    return df, flag


'''技术指标'''


'''行业中性化'''


'''市值中性化'''


